# experiment_recorder.py
import os
import json
import pandas as pd
import matplotlib.pyplot as plt
from datetime import datetime
import torch

class ExperimentRecorder:
    def __init__(self, experiment_name, base_dir="/data/zhouhai/fuxian/实验记录"):
        self.experiment_name = experiment_name
        self.timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        self.experiment_dir = f"{base_dir}/{experiment_name}_{self.timestamp}"
        
        # 创建目录结构
        self.dirs = {
            'models': f"{self.experiment_dir}/models",
            'results': f"{self.experiment_dir}/results",
            'figures': f"{self.experiment_dir}/figures",
            'logs': f"{self.experiment_dir}/logs"
        }
        
        for dir_path in self.dirs.values():
            os.makedirs(dir_path, exist_ok=True)
        
        # 初始化实验日志
        self.log = {
            'metadata': {
                'experiment_name': experiment_name,
                'start_time': self.timestamp,
                'status': 'running'
            },
            'parameters': {},
            'results': {}
        }
    
    def log_parameters(self, params_dict):
        """记录实验参数"""
        self.log['parameters'].update(params_dict)
    
    def log_metrics(self, metrics_dict, epoch=None):
        """记录训练指标"""
        if epoch is not None:
            if 'epochs' not in self.log['results']:
                self.log['results']['epochs'] = {}
            self.log['results']['epochs'][str(epoch)] = metrics_dict
        else:
            self.log['results'].update(metrics_dict)
    
    def save_model(self, model, identifier):
        """保存模型"""
        model_path = f"{self.dirs['models']}/{identifier}.pth"
        torch.save(model.state_dict(), model_path)
        return model_path
    
    def save_figure(self, fig, filename):
        """保存图表"""
        fig_path = f"{self.dirs['figures']}/{filename}.png"
        fig.savefig(fig_path)
        plt.close(fig)
        return fig_path
    
    def save_results(self, df, filename):
        """保存结果表格"""
        csv_path = f"{self.dirs['results']}/{filename}.csv"
        df.to_csv(csv_path, index=False)
        return csv_path
    
    def finalize(self):
        """完成实验记录"""
        self.log['metadata']['end_time'] = datetime.now().strftime("%Y%m%d_%H%M%S")
        self.log['metadata']['status'] = 'completed'
        
        # 保存完整日志
        log_path = f"{self.experiment_dir}/experiment_log.json"
        with open(log_path, 'w') as f:
            json.dump(self.log, f, indent=4)
        
        # 生成报告
        self._generate_report()
        
        print(f"实验记录已完成，保存在: {self.experiment_dir}")
    
    def _generate_report(self):
        """生成实验报告"""
        report = f"""# 实验报告 - {self.experiment_name}

## 基本信息
- 开始时间: {self.log['metadata']['start_time']}
- 结束时间: {self.log['metadata']['end_time']}
- 状态: {self.log['metadata']['status']}

## 实验参数
{self._dict_to_markdown_table(self.log['parameters'])}

## 实验结果
{self._process_results()}
"""
        report_path = f"{self.experiment_dir}/experiment_report.md"
        with open(report_path, 'w') as f:
            f.write(report)
    
    def _process_results(self):
        """处理结果数据"""
        if 'final_metrics' in self.log['results']:
            return self._dict_to_markdown_table(self.log['results']['final_metrics'])
        return "无最终结果记录"
    
    def _dict_to_markdown_table(self, data_dict):
        """将字典转换为Markdown表格"""
        table = "| 参数 | 值 |\n|------|----|\n"
        for key, value in data_dict.items():
            table += f"| {key} | {value} |\n"
        return table